import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.metrics import confusion_matrix
import numpy as np

# 1. 初始化记录容器
train_losses = []
train_accs = []
val_losses = []
val_accs = []

def train_model_with_visualization(model, train_loader, val_loader, criterion, optimizer, num_epochs=10):
    best_acc = 0.0
    
    for epoch in range(num_epochs):
        print(f'Epoch {epoch+1}/{num_epochs}')
        print('-' * 10)
        
        # 训练阶段
        model.train()
        running_loss = 0.0
        running_corrects = 0
        
        for inputs, labels in tqdm(train_loader, desc='Training'):
            inputs = inputs.to(device)
            labels = labels.float().unsqueeze(1).to(device)
            
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            
            running_loss += loss.item() * inputs.size(0)
            preds = (outputs > 0.5).float()
            running_corrects += torch.sum(preds == labels.data)
        
        epoch_loss = running_loss / len(train_loader.dataset)
        epoch_acc = running_corrects.double() / len(train_loader.dataset)
        train_losses.append(epoch_loss)
        train_accs.append(epoch_acc.item())
        
        # 验证阶段
        model.eval()
        val_loss = 0.0
        val_corrects = 0
        all_preds = []
        all_labels = []
        
        with torch.no_grad():
            for inputs, labels in tqdm(val_loader, desc='Validating'):
                inputs = inputs.to(device)
                labels = labels.float().unsqueeze(1).to(device)
                
                outputs = model(inputs)
                loss = criterion(outputs, labels)
                
                val_loss += loss.item() * inputs.size(0)
                preds = (outputs > 0.5).float()
                val_corrects += torch.sum(preds == labels.data)
                
                all_preds.extend(preds.cpu().numpy())
                all_labels.extend(labels.cpu().numpy())
        
        val_loss = val_loss / len(val_loader.dataset)
        val_acc = val_corrects.double() / len(val_loader.dataset)
        val_losses.append(val_loss)
        val_accs.append(val_acc.item())
        
        # 打印epoch结果
        print(f'Train Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}')
        print(f'Val Loss: {val_loss:.4f} Acc: {val_acc:.4f}\n')
        
        # 保存最佳模型
        if val_acc > best_acc:
            best_acc = val_acc
            torch.save(model.state_dict(), 'best_model.pth')
            
            # 保存验证集的预测结果用于可视化
            np.savez('best_predictions.npz', 
                    preds=np.array(all_preds), 
                    labels=np.array(all_labels),
                    images=inputs.cpu().numpy()[:8])  # 保存前8张图片
    
    # ----------------------
    # 训练完成后可视化
    # ----------------------
    plt.figure(figsize=(12, 5))
    
    # 1. Loss曲线
    plt.subplot(1, 2, 1)
    plt.plot(train_losses, label='Train Loss')
    plt.plot(val_losses, label='Val Loss')
    plt.title('Training and Validation Loss')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.legend()
    
    # 2. Accuracy曲线
    plt.subplot(1, 2, 2)
    plt.plot(train_accs, label='Train Acc')
    plt.plot(val_accs, label='Val Acc')
    plt.title('Training and Validation Accuracy')
    plt.xlabel('Epoch')
    plt.ylabel('Accuracy')
    plt.legend()
    
    plt.tight_layout()
    plt.savefig('training_curves.png')
    plt.show()
    
    # 3. 混淆矩阵（使用最佳模型在验证集上的表现）
    plot_confusion_matrix(all_labels, all_preds)
    
    # 4. 示例预测可视化
    visualize_predictions('best_predictions.npz')

def plot_confusion_matrix(true_labels, pred_labels):
    """ 绘制混淆矩阵 """
    cm = confusion_matrix(true_labels, pred_labels)
    plt.figure(figsize=(6, 6))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', 
                xticklabels=['Fake', 'Real'], 
                yticklabels=['Fake', 'Real'])
    plt.title('Confusion Matrix')
    plt.xlabel('Predicted')
    plt.ylabel('True')
    plt.savefig('confusion_matrix.png')
    plt.show()

def visualize_predictions(npz_path):
    """ 可视化预测结果示例 """
    data = np.load(npz_path)
    preds = data['preds']
    labels = data['labels']
    images = data['images']
    
    # 反标准化图像
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    images = images.transpose(0, 2, 3, 1)  # (N, C, H, W) -> (N, H, W, C)
    images = std * images + mean
    images = np.clip(images, 0, 1)
    
    plt.figure(figsize=(12, 6))
    for i in range(min(8, len(images))):
        plt.subplot(2, 4, i+1)
        plt.imshow(images[i])
        plt.title(f'True: {"Real" if labels[i] else "Fake"}\nPred: {"Real" if preds[i] else "Fake"}')
        plt.axis('off')
    plt.tight_layout()
    plt.savefig('prediction_examples.png')
    plt.show()

# 使用带可视化的训练函数
train_model_with_visualization(model, train_loader, val_loader, criterion, optimizer, num_epochs=5)
# 保存完整训练历史
history = {
    'train_loss': train_losses,
    'val_loss': val_losses,
    'train_acc': train_accs,
    'val_acc': val_accs
}
np.save('training_history.npy', history)